# -*- coding: utf-8 -*-
#!/usr/bin/env python
import os
import os.path as osp
import json
from tqdm import tqdm

coco_id_name_map={1: 'person', 2: 'bicycle', 3: 'car', 4: 'motorcycle', 5:
            'airplane',
                               6: 'bus', 7: 'train', 8: 'truck', 9: 'boat', 10:
            'traffic light',
                               11: 'fire hydrant', 13: 'stop sign', 14:
            'parking meter', 15: 'bench',
                               16: 'bird', 17: 'cat', 18: 'dog', 19: 'horse',
            20: 'sheep', 21: 'cow',
                               22: 'elephant', 23: 'bear', 24: 'zebra', 25:
            'giraffe', 27: 'backpack',
                               28: 'umbrella', 31: 'handbag', 32: 'tie', 33:
            'suitcase', 34: 'frisbee',
                               35: 'skis', 36: 'snowboard', 37: 'sports ball',
            38: 'kite', 39: 'baseball bat',
                               40: 'baseball glove', 41: 'skateboard', 42:
            'surfboard', 43: 'tennis racket',
                               44: 'bottle', 46: 'wine glass', 47: 'cup', 48:
            'fork', 49: 'knife', 50: 'spoon',
                               51: 'bowl', 52: 'banana', 53: 'apple', 54:
            'sandwich', 55: 'orange',
                               56: 'broccoli', 57: 'carrot', 58: 'hot dog', 59:
            'pizza', 60: 'donut',
                               61: 'cake', 62: 'chair', 63: 'couch', 64:
            'potted plant', 65: 'bed', 67: 'dining table',
                               70: 'toilet', 72: 'tv', 73: 'laptop', 74:
            'mouse', 75: 'remote', 76: 'keyboard',
                               77: 'cell phone', 78: 'microwave', 79: 'oven',
            80: 'toaster', 81: 'sink',
                               82: 'refrigerator', 84: 'book', 85: 'clock', 86:
            'vase', 87: 'scissors',
                               88: 'teddy bear', 89: 'hair drier', 90:
            'toothbrush'}

def args_parse():
    import argparse
    ap = argparse.ArgumentParser()
    add_arg = ap.add_argument
    add_arg("--json_file", default='data/coco/2017/annotations/instances_train2017.json', help="path to coco_json_file")
    add_arg("--classes_file", default='dataloader/coco_classes_names.txt', help="path to coco_classes_file")
    add_arg("--bboxes_dir", default="bboxes", help="save boxes dir")
    args = ap.parse_args()
    return args

def generate_bboxes(json_info, classes_label, bboxes_dir):
    images_info = json_info["images"]

    for i, instance in enumerate(tqdm(json_info["annotations"])):
        img_id = instance["image_id"]
        bbox = instance["bbox"]
        category = coco_id_name_map[instance["category_id"]]
        if category not in classes_label:
            continue
        for image_info in images_info:
            if image_info['id'] != img_id:
                continue
            img_name = image_info["file_name"]
            height = image_info["height"]
            width = image_info["width"]
            break
        x1, y1, w, h = bbox
        x2, y2 = x1 + w, y1 + h
        x1, x2  = x1 / width, x2 / width
        y1, y2 = y1 / height, y2 / height
        instance_bbox = [x1, y1, x2, y2]
        str_bbox = ' '.join(map(str, instance_bbox))
        line_str = ' '.join([category, str_bbox]) + '\n'
        with open(osp.join(bboxes_dir, '.'.join([img_name.split('.')[0],\
                                                 'txt'])), 'a') as f:
            f.write(line_str)

def load_json(json_file):
    json_info = None
    with open(json_file, 'r') as f:
        json_info = json.load(f)
    #with open('trans_coco.json', 'w') as f_out:
    #    json.dump(json_info, f_out, ensure_ascii=False, indent=2)
    return json_info

def dump_json(json_file):
    json_info = None
    with open(json_file, 'r') as f:
        json_info = json.load(f)

    with open("trans_coco.json", "w") as f_out:
        json.dump(json_info, f_out, ensure_ascii=False, indent=2)

def load_classes_names(classes_file):
    classes_label = {}
    with open(classes_file, 'r') as f:
        classes_name = f.readlines()
        for i, name in enumerate(classes_name):
            classes_label[name.strip()] = i
    return classes_label

if __name__=="__main__":
    args = args_parse()
    # print(args.json_file)
    json_info = load_json(args.json_file)
    classes_label = load_classes_names(args.classes_file)

    generate_bboxes(json_info, classes_label, args.bboxes_dir)
